"""
PaddleOCR识别服务模块
支持批量识别和多线程处理，优化中英文混合文本识别
"""

import time
import threading
from typing import List, Dict, Optional, Tuple, Any
from dataclasses import dataclass
from concurrent.futures import ThreadPoolExecutor, Future
import numpy as np
import cv2
from paddleocr import PaddleOCR
import queue

from ..core import config, ocr_logger
from ..processing import CroppedRegion


@dataclass
class OCRResult:
    """OCR识别结果"""
    region_id: str
    text: str
    confidence: float
    bbox: List[List[int]]  # 文本行的边界框坐标
    timestamp: float
    processing_time: float


@dataclass
class BatchOCRResult:
    """批量OCR识别结果"""
    results: List[OCRResult]
    batch_size: int
    total_processing_time: float
    average_confidence: float


class PaddleOCRService:
    """PaddleOCR识别服务"""
    
    def __init__(self):
        self.use_gpu = config.ocr.use_gpu
        self.use_angle_cls = config.ocr.use_angle_cls
        self.use_space_char = config.ocr.use_space_char
        self.lang = config.ocr.lang
        self.batch_size = config.ocr.batch_size
        self.confidence_threshold = config.ocr.confidence_threshold
        self.max_text_length = config.ocr.max_text_length
        
        # OCR引擎
        self.ocr_engine = None
        self.ocr_lock = threading.Lock()
        
        # 批处理队列
        self.batch_queue = queue.Queue()
        self.batch_results = {}
        self.batch_lock = threading.Lock()
        
        # 线程池
        self.executor = ThreadPoolExecutor(max_workers=4)
        
        # 统计信息
        self.stats = {
            "total_processed": 0,
            "total_text_length": 0,
            "average_confidence": 0.0,
            "processing_time": 0.0
        }
        
        self._initialize_ocr()
        ocr_logger.info("PaddleOCR服务初始化完成")
    
    def _initialize_ocr(self):
        """初始化OCR引擎"""
        try:
            ocr_logger.info("正在初始化PaddleOCR引擎...")
            
            self.ocr_engine = PaddleOCR(
                use_angle_cls=self.use_angle_cls,
                lang=self.lang,
                use_gpu=self.use_gpu,
                use_space_char=self.use_space_char,
                show_log=False
            )
            
            ocr_logger.info(f"PaddleOCR引擎初始化成功，语言: {self.lang}, GPU: {self.use_gpu}")
            
        except Exception as e:
            ocr_logger.error(f"PaddleOCR引擎初始化失败: {e}")
            raise
    
    def recognize_single(self, region: CroppedRegion) -> OCRResult:
        """识别单个区域"""
        start_time = time.time()
        
        try:
            # 预处理图像
            processed_image = self._preprocess_image(region.image)
            
            # OCR识别
            with self.ocr_lock:
                results = self.ocr_engine.ocr(processed_image, cls=self.use_angle_cls)
            
            # 解析结果
            text_lines = []
            confidences = []
            bboxes = []
            
            if results and results[0]:
                for line in results[0]:
                    if line and len(line) >= 2:
                        bbox, (text, confidence) = line
                        
                        # 过滤低置信度结果
                        if confidence >= self.confidence_threshold:
                            text_lines.append(text)
                            confidences.append(confidence)
                            bboxes.append(bbox)
            
            # 合并文本
            combined_text = ' '.join(text_lines) if text_lines else ""
            average_confidence = sum(confidences) / len(confidences) if confidences else 0.0
            
            # 限制文本长度
            if len(combined_text) > self.max_text_length:
                combined_text = combined_text[:self.max_text_length] + "..."
            
            processing_time = time.time() - start_time
            
            # 更新统计信息
            self._update_stats(combined_text, average_confidence, processing_time)
            
            result = OCRResult(
                region_id=region.region_id,
                text=combined_text,
                confidence=average_confidence,
                bbox=bboxes,
                timestamp=time.time(),
                processing_time=processing_time
            )
            
            ocr_logger.debug(
                f"OCR识别完成: {region.region_id}, "
                f"文本长度: {len(combined_text)}, "
                f"置信度: {average_confidence:.3f}, "
                f"耗时: {processing_time:.3f}s"
            )
            
            return result
            
        except Exception as e:
            ocr_logger.error(f"OCR识别失败 {region.region_id}: {e}")
            return OCRResult(
                region_id=region.region_id,
                text="",
                confidence=0.0,
                bbox=[],
                timestamp=time.time(),
                processing_time=time.time() - start_time
            )
    
    def recognize_batch(self, regions: List[CroppedRegion]) -> BatchOCRResult:
        """批量识别多个区域"""
        start_time = time.time()
        
        if not regions:
            return BatchOCRResult(
                results=[],
                batch_size=0,
                total_processing_time=0.0,
                average_confidence=0.0
            )
        
        ocr_logger.info(f"开始批量OCR识别，批次大小: {len(regions)}")
        
        # 分批处理
        batch_results = []
        batch_size = min(self.batch_size, len(regions))
        
        for i in range(0, len(regions), batch_size):
            batch = regions[i:i + batch_size]
            
            # 并行处理当前批次
            futures = []
            for region in batch:
                future = self.executor.submit(self.recognize_single, region)
                futures.append(future)
            
            # 收集结果
            for future in futures:
                try:
                    result = future.result(timeout=30)  # 30秒超时
                    batch_results.append(result)
                except Exception as e:
                    ocr_logger.error(f"批处理中的任务失败: {e}")
        
        total_processing_time = time.time() - start_time
        
        # 计算平均置信度
        confidences = [r.confidence for r in batch_results if r.confidence > 0]
        average_confidence = sum(confidences) / len(confidences) if confidences else 0.0
        
        result = BatchOCRResult(
            results=batch_results,
            batch_size=len(regions),
            total_processing_time=total_processing_time,
            average_confidence=average_confidence
        )
        
        ocr_logger.info(
            f"批量OCR识别完成，处理 {len(batch_results)}/{len(regions)} 个区域，"
            f"平均置信度: {average_confidence:.3f}, "
            f"总耗时: {total_processing_time:.3f}s"
        )
        
        return result
    
    def _preprocess_image(self, image: np.ndarray) -> np.ndarray:
        """预处理图像以提高OCR识别效果"""
        try:
            # 确保图像是BGR格式
            if len(image.shape) == 3 and image.shape[2] == 3:
                processed = image.copy()
            else:
                processed = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
            
            # 图像增强
            processed = self._enhance_image(processed)
            
            return processed
            
        except Exception as e:
            ocr_logger.error(f"图像预处理失败: {e}")
            return image
    
    def _enhance_image(self, image: np.ndarray) -> np.ndarray:
        """增强图像质量"""
        try:
            # 转换为灰度图
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            
            # 自适应直方图均衡化
            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
            enhanced = clahe.apply(gray)
            
            # 双边滤波去噪
            enhanced = cv2.bilateralFilter(enhanced, 9, 75, 75)
            
            # 锐化
            kernel = np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]])
            enhanced = cv2.filter2D(enhanced, -1, kernel)
            
            # 转回BGR格式
            enhanced = cv2.cvtColor(enhanced, cv2.COLOR_GRAY2BGR)
            
            return enhanced
            
        except Exception as e:
            ocr_logger.error(f"图像增强失败: {e}")
            return image
    
    def _update_stats(self, text: str, confidence: float, processing_time: float):
        """更新统计信息"""
        self.stats["total_processed"] += 1
        self.stats["total_text_length"] += len(text)
        self.stats["processing_time"] += processing_time
        
        # 计算平均置信度
        total_confidence = self.stats["average_confidence"] * (self.stats["total_processed"] - 1)
        self.stats["average_confidence"] = (total_confidence + confidence) / self.stats["total_processed"]
    
    def get_stats(self) -> Dict[str, Any]:
        """获取统计信息"""
        stats = self.stats.copy()
        if stats["total_processed"] > 0:
            stats["average_processing_time"] = stats["processing_time"] / stats["total_processed"]
            stats["average_text_length"] = stats["total_text_length"] / stats["total_processed"]
        else:
            stats["average_processing_time"] = 0.0
            stats["average_text_length"] = 0.0
        
        return stats
    
    def reset_stats(self):
        """重置统计信息"""
        self.stats = {
            "total_processed": 0,
            "total_text_length": 0,
            "average_confidence": 0.0,
            "processing_time": 0.0
        }
        ocr_logger.info("OCR统计信息已重置")
    
    def cleanup(self):
        """清理资源"""
        ocr_logger.info("清理OCR服务资源")
        
        # 关闭线程池
        self.executor.shutdown(wait=True)
        
        # 清理队列
        while not self.batch_queue.empty():
            try:
                self.batch_queue.get_nowait()
            except queue.Empty:
                break
        
        ocr_logger.info("OCR服务资源清理完成")


class OCRBatchProcessor:
    """OCR批处理器"""
    
    def __init__(self, ocr_service: PaddleOCRService):
        self.ocr_service = ocr_service
        self.batch_queue = queue.Queue()
        self.result_callbacks = {}
        self.is_running = False
        self.worker_thread = None
        
        ocr_logger.info("OCR批处理器初始化完成")
    
    def start(self):
        """启动批处理器"""
        if self.is_running:
            return
        
        self.is_running = True
        self.worker_thread = threading.Thread(target=self._worker_loop, daemon=True)
        self.worker_thread.start()
        
        ocr_logger.info("OCR批处理器已启动")
    
    def stop(self):
        """停止批处理器"""
        if not self.is_running:
            return
        
        self.is_running = False
        if self.worker_thread and self.worker_thread.is_alive():
            self.worker_thread.join(timeout=5)
        
        ocr_logger.info("OCR批处理器已停止")
    
    def submit_batch(self, regions: List[CroppedRegion], 
                    callback: Optional[callable] = None) -> str:
        """提交批处理任务"""
        batch_id = f"batch_{int(time.time() * 1000)}"
        
        if callback:
            self.result_callbacks[batch_id] = callback
        
        self.batch_queue.put((batch_id, regions))
        
        ocr_logger.debug(f"提交OCR批处理任务: {batch_id}, 区域数: {len(regions)}")
        return batch_id
    
    def _worker_loop(self):
        """工作线程循环"""
        ocr_logger.info("OCR批处理工作线程启动")
        
        while self.is_running:
            try:
                # 获取批处理任务
                batch_id, regions = self.batch_queue.get(timeout=1)
                
                # 执行批处理
                result = self.ocr_service.recognize_batch(regions)
                
                # 调用回调函数
                if batch_id in self.result_callbacks:
                    try:
                        self.result_callbacks[batch_id](batch_id, result)
                    except Exception as e:
                        ocr_logger.error(f"批处理回调失败 {batch_id}: {e}")
                    finally:
                        del self.result_callbacks[batch_id]
                
            except queue.Empty:
                continue
            except Exception as e:
                ocr_logger.error(f"批处理工作线程异常: {e}")
        
        ocr_logger.info("OCR批处理工作线程结束")
